from transformers import RobertaForSequenceClassification, RobertaTokenizer
import torch
import torch.nn.functional as F

# 本地模型路径
model_path = "/e/Resources/LLM/huggingface/roberta-base-go_emotions"
# 加载本地模型和分词器
model = RobertaForSequenceClassification.from_pretrained(model_path)
tokenizer = RobertaTokenizer.from_pretrained(model_path)


#如果从模型读取（从huggingface下载文件）
#model_name = "roberta-base-go_emotions"
#model = RobertaForSequenceClassification.from_pretrained(model_name)
#tokenizer = RobertaTokenizer.from_pretrained(model_name)

# 输入文本
text = "I am so happy today!"

# 对文本进行编码
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)

# gpu支持
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
inputs = {k: v.to(device) for k, v in inputs.items()}

# 禁用梯度计算
with torch.no_grad():
    outputs = model(**inputs)

# 获取预测结果
logits = outputs.logits

# 将 logits 转换为概率
probs = F.softmax(logits, dim=-1)

# 获取概率最高的类别
predicted_class = torch.argmax(probs, dim=-1).item()

# 情感类别列表
emotion_labels = [
    "admiration", "amusement", "anger", "annoyance", "approval", "caring", "confusion", "curiosity", 
    "desire", "disappointment", "disapproval", "disgust", "embarrassment", "excitement", "fear", 
    "gratitude", "grief", "joy", "love", "nervousness", "optimism", "pride", "realization", 
    "relief", "remorse", "sadness", "surprise", "neutral"
]

# 打印预测的情感类别
print(f"Predicted emotion: {emotion_labels[predicted_class]}")